Skip to content

[None][fix] Bypass the sync transfer-admission budget on the async disagg path#16834

Draft
Shixiaowei02 wants to merge 2 commits into
NVIDIA:mainfrom
Shixiaowei02:user/xiaoweis/fix-disagg-async-transfer-throttle
Draft

[None][fix] Bypass the sync transfer-admission budget on the async disagg path#16834
Shixiaowei02 wants to merge 2 commits into
NVIDIA:mainfrom
Shixiaowei02:user/xiaoweis/fix-disagg-async-transfer-throttle

Conversation

@Shixiaowei02

@Shixiaowei02 Shixiaowei02 commented Jul 24, 2026 •
edited by coderabbitai Bot
Loading

Copy link
Copy Markdown
Collaborator

Description

This pull request improves the robustness and concurrency handling of the disaggregated KV-cache transfer logic in PyExecutor, and enhances test coverage for these scenarios. The main changes include allowing asynchronous transfer paths to bypass per-iteration transfer budgets, replacing a hard assertion with a warning to handle transient request overshoots, and updating related unit tests to cover these behaviors.

Concurrency and transfer admission improvements:

  • Updated _apply_disagg_transfer_admission in py_executor.py to allow asynchronous disaggregated KV-cache transfers to bypass the per-iteration transfer budget, preventing unnecessary throttling and timeouts at high concurrency.
  • In start_thread of worker.py, added logic to prevent thread restarts after engine event loop crashes, surfacing the actual error instead of masking it.

Robustness improvements:

  • Replaced a hard assertion in _pad_attention_dp_dummy_request with a warning when active_requests temporarily exceeds expected_num_active_requests, preventing crashes during transient overshoots caused by lingering disagg transfer-error requests.

Testing and coverage:

  • Added and updated unit tests in test_py_executor.py to:
    • Explicitly test that the transfer budget only gates the synchronous path and is bypassed in async mode.
    • Ensure that the system tolerates transient overshoots of active requests without raising exceptions.
    • Pin modes in existing tests to exercise both sync and async code paths, ensuring correct gating logic. [1] [2]

Test Coverage

PR Checklist

Please review the following before submitting your PR:

  • PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

  • PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

  • Test cases are provided for new code paths (see test instructions)

  • If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

  • Any new dependencies have been scanned for license and vulnerabilities

  • CODEOWNERS updated if ownership changes

  • Documentation updated as needed

  • Update tava architecture diagram if there is a significant design change in PR.

  • The reviewers assigned automatically/manually are appropriate for the PR.

  • Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

Dev Engineer Review

  • Async disaggregated KV transfers now bypass the synchronous transfer-admission budget, avoiding unnecessary throttling while retaining synchronous admission control.
  • Finished response threads are no longer restarted; captured event-loop errors are re-raised, preventing misleading restart failures and potential MPI hangs.
  • Active-request overshoots now emit a warning instead of crashing on an assertion, tolerating transient cleanup states.
  • No configuration or test-list files were changed.

QA Engineer Review

Test functions added or modified:

  • test_apply_reverts_deferred_v2_allocations
  • test_apply_missing_v2_flag_defaults_to_non_v2
  • test_async_mode_bypasses_transfer_budget
  • test_pad_dummy_tolerates_active_request_overshoot

These are unit tests in tests/unittest/_torch/executor/test_py_executor.py; no corresponding test-db/ or qa/ test-list entries were added or modified.

Verdict: sufficient

@Shixiaowei02
Shixiaowei02 requested review from a team as code owners July 24, 2026 08:51
@Shixiaowei02
Shixiaowei02 marked this pull request as draft July 24, 2026 08:57
@coderabbitai

coderabbitai Bot commented Jul 24, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 590e7a51-1c05-4ebd-b87e-26607e813469

📥 Commits

Reviewing files that changed from the base of the PR and between 75b39d4 and 84fa5a5.

📒 Files selected for processing (3)
  • tensorrt_llm/_torch/pyexecutor/py_executor.py
  • tensorrt_llm/executor/worker.py
  • tests/unittest/_torch/executor/test_py_executor.py

Walkthrough

Changes

Transfer admission behavior

Layer / File(s) Summary
Async transfer admission bypass and coverage
tensorrt_llm/_torch/pyexecutor/py_executor.py, tests/unittest/_torch/executor/test_py_executor.py		 Async generation transfers bypass synchronous budget admission. Existing budgeting tests force synchronous overlap behavior, while async-mode coverage verifies candidates remain admitted without rollback.

Executor runtime safeguards

Layer / File(s) Summary
Active-request padding tolerance
tensorrt_llm/_torch/pyexecutor/py_executor.py, tests/unittest/_torch/executor/test_py_executor.py		 Transient active-request overshoot logs a warning instead of raising an assertion, with regression coverage.
Response-thread lifecycle guards
tensorrt_llm/executor/worker.py		 Response-thread startup avoids invalid restarts and re-raises captured event-loop errors.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Suggested reviewers: lancelly, chienchunhung

🚥 Pre-merge checks | ✅ 5
✅ Passed checks (5 passed)
Check name Status Explanation
Title check ✅ Passed The title is concise and accurately describes the main async transfer-budget bypass change.
Description check ✅ Passed The description clearly explains the issue, fix, and tests; the Test Coverage section is present but not separately filled.
Docstring Coverage ✅ Passed No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

Comment @coderabbitai help to get the list of available commands.

@Shixiaowei02
[None][fix] Surface engine event-loop crash instead of a thread-resta…
5f0f52d 
…rt cascade

start_thread() restarted the already-finished await-response thread, raising
"threads can only be started once" and masking the real cause: the engine event
loop crashed (AwaitResponseHelper broadcasts _event_loop_error and returns, so
the thread exits). The misleading RuntimeError then cascaded into a peer
MPI-collective hang across ranks. Detect the finished thread and re-raise the
stashed engine event-loop error instead.

Signed-off-by: Shixiaowei02 <39303645+Shixiaowei02@users.noreply.github.com>
@Shixiaowei02
[None][fix] Don't throttle async disagg gen KV transfers by max_token…
6d1195c 
…s_in_buffer

The DisaggTransferAdmissionController derives its in-flight transfer budget from
max_tokens_in_buffer (a gather-buffer sizing knob), which for an 8k / high-
concurrency config is only ~2 concurrent transfers per rank. That bound is meant
for the SYNCHRONOUS transfer path (limit blocking transfers per executor
iteration); on the ASYNC path transfers stay in flight and don't block, so
applying it there serializes the ctx->gen handoff and backs requests up past
kv_transfer_timeout_ms, failing them with "Disagg KV cache transfer error".
Bypass the budget on the async path (already capacity-scheduled, so bounded by
max_batch_size, matching v1/CPP behavior).

Also harden _pad_attention_dp_dummy_request: a transient overshoot where
KV-transfer-error requests briefly linger in active_requests (len >
expected_num_active_requests) tripped a hard assert that crashed the gen loop on
all ADP ranks. Warn and continue; the padding decision keys on the schedulable
count.

Signed-off-by: Shixiaowei02 <39303645+Shixiaowei02@users.noreply.github.com>
@Shixiaowei02
Shixiaowei02 force-pushed the user/xiaoweis/fix-disagg-async-transfer-throttle branch from 84fa5a5 to 6d1195c Compare July 24, 2026 10:03
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@cascade812 cascade812 Awaiting requested review from cascade812 cascade812 is a code owner automatically assigned from NVIDIA/trt-llm-runtime-devs
@tongyuantongyu tongyuantongyu Awaiting requested review from tongyuantongyu tongyuantongyu is a code owner automatically assigned from NVIDIA/trt-llm-runtime-devs
@chuangz0 chuangz0 Awaiting requested review from chuangz0 chuangz0 is a code owner automatically assigned from NVIDIA/trt-llm-disagg-devs
@bo-nv bo-nv Awaiting requested review from bo-nv bo-nv is a code owner automatically assigned from NVIDIA/trt-llm-disagg-devs
@chienchunhung chienchunhung Awaiting requested review from chienchunhung
@asfiyab-nvidia asfiyab-nvidia Awaiting requested review from asfiyab-nvidia
@nv-xtf nv-xtf Awaiting requested review from nv-xtf

At least 1 approving review is required to merge this pull request.

Assignees

@Shixiaowei02 Shixiaowei02

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@Shixiaowei02